from __future__ import annotations

from datetime import datetime
from typing import Optional

import pandas as pd

from core.database import db_manager, DatabaseManager
from dao.models import ContractInfo


class ContractInfoDao:
    def __init__(self, db: DatabaseManager | None = None):
        self.db = db or db_manager

    def get_all_contract_info(self) -> pd.DataFrame:
        """获取所有合约信息"""
        session = self.db.get_session()
        try:
            rows = session.query(ContractInfo).all()
            if not rows:
                return pd.DataFrame(columns=['vnpy_underlying_symbol', 'contract_size', 'exchange', 'name'])
            data = [
                {
                    'vnpy_underlying_symbol': r.vnpy_underlying_symbol,
                    'contract_size': r.contract_size,
                    'exchange': r.exchange,
                    'name': r.name,
                }
                for r in rows
            ]
            return pd.DataFrame(data)
        finally:
            session.close()

    def get_contract_info(self, symbol: str, exchange: str) -> Optional[float]:
        """获取单个合约乘数"""
        session = self.db.get_session()
        try:
            row = (
                session.query(ContractInfo)
                .filter(ContractInfo.vnpy_underlying_symbol == symbol, ContractInfo.exchange == exchange)
                .first()
            )
            return row.contract_size if row else None
        finally:
            session.close()

    def save_contract_info(self, df: pd.DataFrame):
        """保存合约信息到数据库"""
        if df.empty:
            return
        session = self.db.get_session()
        try:
            now = datetime.utcnow()
            for _, row in df.iterrows():
                sym = row['vnpy_underlying_symbol']
                item: ContractInfo | None = (
                    session.query(ContractInfo).filter(ContractInfo.vnpy_underlying_symbol == sym).first()
                )
                if item:
                    item.contract_size = float(row['contract_size'])
                    item.exchange = str(row['exchange'])
                    item.name = str(row['name'])
                    item.update_time = now
                else:
                    item = ContractInfo(
                        vnpy_underlying_symbol=str(row['vnpy_underlying_symbol']),
                        contract_size=float(row['contract_size']),
                        exchange=str(row['exchange']),
                        name=str(row['name']),
                        create_time=now,
                        update_time=now,
                    )
                    session.add(item)
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

